#include <stdlib.h>
#include <algorithm>
#include <chrono>
#include <iostream>
#include <vector>
#include <memory>
#include <functional>

#include <rclcpp/rclcpp.hpp>
#include <std_msgs/msg/string.hpp>

#include <geometry_msgs/msg/pose_array.hpp>
#include <geometry_msgs/msg/pose_stamped.hpp>
#include <geometry_msgs/msg/pose_with_covariance_stamped.hpp>
#include <sensor_msgs/msg/point_cloud.hpp>

#include "vehicle_msgs/msg/control_signal.hpp"
// #include "vehicle_msgs/decoder.h"  // Temporarily disabled for ROS2 compatibility

#include "phy_simulator/basics.h"
#include "phy_simulator/phy_simulator.h"
#include "phy_simulator/ros_adapter.h"
#include "phy_simulator/visualizer.h"

using namespace phy_simulator;

DECLARE_BACKWARD;
const double simulation_rate = 500.0;
const double gt_msg_rate = 100.0;
const double gt_static_msg_rate = 10.0;
const double visualization_msg_rate = 20.0;

class PhySimulatorPlanningNode : public rclcpp::Node
{
public:
  PhySimulatorPlanningNode() : Node("phy_simulator_planning_node")
  {
    // Declare parameters
    this->declare_parameter("vehicle_info_path", "");
    this->declare_parameter("map_path", "");
    this->declare_parameter("lane_net_path", "");
    
    // Get parameters
    std::string vehicle_info_path = this->get_parameter("vehicle_info_path").as_string();
    std::string map_path = this->get_parameter("map_path").as_string();
    std::string lane_net_path = this->get_parameter("lane_net_path").as_string();
    
    if (vehicle_info_path.empty() || map_path.empty() || lane_net_path.empty()) {
      RCLCPP_ERROR(this->get_logger(), "Failed to get required parameters");
      return;
    }
    
    RCLCPP_INFO(this->get_logger(), "[PhySimulation] Initialization finished, waiting for callback");
  }
  
  void initialize() {
    // Initialize physics simulation
    auto param_values = this->get_parameters({"vehicle_info_path", "map_path", "lane_net_path"});
    std::string vehicle_info_path = param_values[0].as_string();
    std::string map_path = param_values[1].as_string();
    std::string lane_net_path = param_values[2].as_string();
    
    if (!vehicle_info_path.empty() && !map_path.empty() && !lane_net_path.empty()) {
      try {
        phy_sim_ = std::make_unique<PhySimulation>(vehicle_info_path, map_path, lane_net_path);
        RCLCPP_INFO(this->get_logger(), "Physics simulation initialized successfully");
      } catch (const std::exception& e) {
        RCLCPP_ERROR(this->get_logger(), "Failed to initialize physics simulation: %s", e.what());
        phy_sim_.reset();
      }
    } else {
      RCLCPP_WARN(this->get_logger(), "Physics simulation not initialized due to missing parameters");
    }
    
    // Initialize adapters
    ros_adapter_ = std::make_unique<RosAdapter>(shared_from_this());
    if (phy_sim_) {
      ros_adapter_->set_phy_sim(phy_sim_.get());
    }
    
    visualizer_ = std::make_unique<Visualizer>(shared_from_this());
    if (phy_sim_) {
      visualizer_->set_phy_sim(phy_sim_.get());
    }
    
    // Setup vehicle subscriptions only if physics simulation is initialized
    if (phy_sim_) {
      auto vehicle_ids = phy_sim_->vehicle_ids();
      int num_vehicles = static_cast<int>(vehicle_ids.size());
      ctrl_subscribers_.resize(num_vehicles);
      
      for (int i = 0; i < num_vehicles; i++) {
        auto vehicle_id = vehicle_ids[i];
        std::string topic_name = std::string("/ctrl/agent_") + std::to_string(vehicle_id);
        RCLCPP_INFO(this->get_logger(), "Subscribing to %s", topic_name.c_str());
        
        ctrl_subscribers_[i] = this->create_subscription<vehicle_msgs::msg::ControlSignal>(
          topic_name, 10,
          [this, vehicle_id](const vehicle_msgs::msg::ControlSignal::SharedPtr msg) {
            this->CtrlSignalCallback(msg, vehicle_id);
          });
      }
      
      // Initialize signal set
      for (auto& vehicle_id : vehicle_ids) {
        common::VehicleControlSignal default_signal;
        signal_set_.signal_set.insert(std::pair<int, common::VehicleControlSignal>(
            vehicle_id, default_signal));
      }
    }
    
    // Setup other subscriptions
    initial_pose_sub_ = this->create_subscription<geometry_msgs::msg::PoseWithCovarianceStamped>(
      "/initialpose", 10,
      std::bind(&PhySimulatorPlanningNode::InitialPoseCallback, this, std::placeholders::_1));
    
    goal_pose_sub_ = this->create_subscription<geometry_msgs::msg::PoseStamped>(
      "/move_base_simple/goal", 10,
      std::bind(&PhySimulatorPlanningNode::NavGoalCallback, this, std::placeholders::_1));
    
    // Setup timers
    simulation_timer_ = this->create_wall_timer(
      std::chrono::microseconds(static_cast<int>(1e6 / simulation_rate)),
      std::bind(&PhySimulatorPlanningNode::SimulationTimerCallback, this));
    
    gt_pub_timer_ = this->create_wall_timer(
      std::chrono::microseconds(static_cast<int>(1e6 / gt_msg_rate)),
      std::bind(&PhySimulatorPlanningNode::GtPubTimerCallback, this));
    
    gt_static_pub_timer_ = this->create_wall_timer(
      std::chrono::microseconds(static_cast<int>(1e6 / gt_static_msg_rate)),
      std::bind(&PhySimulatorPlanningNode::GtStaticPubTimerCallback, this));
    
    vis_pub_timer_ = this->create_wall_timer(
      std::chrono::microseconds(static_cast<int>(1e6 / visualization_msg_rate)),
      std::bind(&PhySimulatorPlanningNode::VisPubTimerCallback, this));
  }

private:
  void CtrlSignalCallback(const vehicle_msgs::msg::ControlSignal::SharedPtr msg, int index)
  {
    common::VehicleControlSignal ctrl;
    // Note: decoder needs to be updated for ROS2 - using manual conversion for now
    ctrl.acc = msg->acc;
    ctrl.steer_rate = msg->steer_rate;
    ctrl.is_openloop = msg->is_openloop;
    
    // Convert state manually for now
    ctrl.state.vec_position[0] = msg->state.vec_position.x;
    ctrl.state.vec_position[1] = msg->state.vec_position.y;
    ctrl.state.angle = msg->state.angle;
    ctrl.state.curvature = msg->state.curvature;
    ctrl.state.velocity = msg->state.velocity;
    ctrl.state.acceleration = msg->state.acceleration;
    ctrl.state.steer = msg->state.steer;
    
    signal_set_.signal_set[index] = ctrl;
    RCLCPP_DEBUG(this->get_logger(), "Received control signal for vehicle %d", index);
  }
  
  void InitialPoseCallback(const geometry_msgs::msg::PoseWithCovarianceStamped::SharedPtr msg)
  {
    // Convert ROS2 pose to internal representation
    initial_state_[0] = msg->pose.pose.position.x;
    initial_state_[1] = msg->pose.pose.position.y;
    
    // Convert quaternion to yaw angle
    auto q = msg->pose.pose.orientation;
    initial_state_[2] = atan2(2.0 * (q.w * q.z + q.x * q.y), 1.0 - 2.0 * (q.y * q.y + q.z * q.z));
    
    flag_rcv_initial_state_ = true;
    RCLCPP_INFO(this->get_logger(), "Received initial pose: [%.2f, %.2f, %.2f]", 
                initial_state_[0], initial_state_[1], initial_state_[2]);
  }
  
  void NavGoalCallback(const geometry_msgs::msg::PoseStamped::SharedPtr msg)
  {
    // Convert ROS2 pose to internal representation
    goal_state_[0] = msg->pose.position.x;
    goal_state_[1] = msg->pose.position.y;
    
    // Convert quaternion to yaw angle
    auto q = msg->pose.orientation;
    goal_state_[2] = atan2(2.0 * (q.w * q.z + q.x * q.y), 1.0 - 2.0 * (q.y * q.y + q.z * q.z));
    
    flag_rcv_goal_state_ = true;
    RCLCPP_INFO(this->get_logger(), "Received navigation goal: [%.2f, %.2f, %.2f]", 
                goal_state_[0], goal_state_[1], goal_state_[2]);
  }
  
  void SimulationTimerCallback()
  {
    if (phy_sim_) {
      phy_sim_->UpdateSimulatorUsingSignalSet(signal_set_, 1.0 / simulation_rate);
    }
    RCLCPP_DEBUG(this->get_logger(), "Simulation timer callback");
  }
  
  void GtPubTimerCallback()
  {
    if (!ros_adapter_) {
      return; // Skip if adapter is not initialized
    }
    
    auto now = this->get_clock()->now();
    ros_adapter_->PublishDynamicDataWithStamp(now);
  }
  
  void GtStaticPubTimerCallback()
  {
    if (!ros_adapter_) {
      return; // Skip if adapter is not initialized
    }
    
    auto now = this->get_clock()->now();
    ros_adapter_->PublishStaticDataWithStamp(now);
  }
  
  void VisPubTimerCallback()
  {
    if (!visualizer_) {
      return; // Skip if visualizer is not initialized
    }
    
    auto now = this->get_clock()->now();
    visualizer_->VisualizeDataWithStamp(now);
  }
  
  // Member variables
  std::unique_ptr<PhySimulation> phy_sim_;
  std::unique_ptr<RosAdapter> ros_adapter_;
  std::unique_ptr<Visualizer> visualizer_;
  
  common::VehicleControlSignalSet signal_set_;
  std::vector<rclcpp::Subscription<vehicle_msgs::msg::ControlSignal>::SharedPtr> ctrl_subscribers_;
  
  rclcpp::Subscription<geometry_msgs::msg::PoseWithCovarianceStamped>::SharedPtr initial_pose_sub_;
  rclcpp::Subscription<geometry_msgs::msg::PoseStamped>::SharedPtr goal_pose_sub_;
  
  rclcpp::TimerBase::SharedPtr simulation_timer_;
  rclcpp::TimerBase::SharedPtr gt_pub_timer_;
  rclcpp::TimerBase::SharedPtr gt_static_pub_timer_;
  rclcpp::TimerBase::SharedPtr vis_pub_timer_;
  
  Vec3f initial_state_{0, 0, 0};
  bool flag_rcv_initial_state_ = false;
  
  Vec3f goal_state_{0, 0, 0};
  bool flag_rcv_goal_state_ = false;
};

int main(int argc, char ** argv)
{
  rclcpp::init(argc, argv);
  
  auto node = std::make_shared<PhySimulatorPlanningNode>();
  
  // Initialize after the shared_ptr is created
  node->initialize();
  
  RCLCPP_INFO(node->get_logger(), "Starting phy_simulator planning node");
  
  rclcpp::spin(node);
  
  rclcpp::shutdown();
  return 0;
}
